package ecc;

/*
 * Elliptic curve signature engine.
 *
 * ================================
 *
 * PSEC-1 encryption scheme.
 *
 * Notation:
 *
 * G: base point on the curve.
 * n: prime order of the base point
 * m: message representative (0 <= m < n)
 * x: private key (1 <= x < n)
 * Y: public key (Y = x*G)
 *
 * Signature generation:
 *
 * k = random()                         (condition: 1 <= k < n)
 * W = k*G
 * t = abscissa(W) mod n
 * r = t                                (condition: r != 0)
 * s = k^(-1)*(m + r*x) mod n           (condition: s != 0)
 *
 * The signature is the pair (r, s).
 *
 * Signature verification:
 *
 * z = s^(-1) mod n                     (condition: 1 <= s < n)
 * u = z*m mod n
 * v = z*r mod n                        (condition: 1 <= r < n)
 * W = u*G + v*Y = s^(-1)*(m + r*x)*G = k*G
 * t = abscissa(W) mod n
 *
 * The signature is accepted <=> t = r.
 *
 * ================================
 *
 * Nyberg-Rueppel signature scheme.
 *
 * Notation:
 *
 * G: base point on the curve.
 * n: prime order of the base point
 * m: message representative (0 <= m < n)
 * x: private key (1 <= x < n)
 * Y: public key (Y = x*G)
 *
 * Signature generation:
 *
 * k = random()                         (condition: 1 <= k < n)
 * W = k*G
 * t = abscissa(W) mod n
 * r = (m + t) mod n                    (condition: r != 0)
 * s = (k - r*x) mod n
 *
 * The signature is the pair (r, s).
 *
 * Signature verification:
 *
 * W = s*G + r*Y = (s + r*x)*G = k*G    (condition: 1 <= r < n, 0 <= s < n)
 * t = abscissa(W) mod n
 * z = (r - t) mod n
 *
 * The signature is accepted <=> z = m.
 *
 * ================================
 *
 * @author  Paulo S. L. M. Barreto
 */
import java.math.BigInteger;
import java.security.GeneralSecurityException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.MessageDigest;
import java.util.Arrays;
import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.ShortBufferException;
import javax.crypto.spec.SecretKeySpec;

public class EEncrypt {

	private static final int STATE_NOT_INITIALIZED = 0;
	private static final int STATE_ENCRYPTING = 1;
	private static final int STATE_DECRYPTING = 2;
	private int state = STATE_NOT_INITIALIZED;

	private static final int CIPHER_KEY_LENGTH = 16; //AES key

	protected MessageDigest md = null;
	private ECurve E = null;
	private EPoint C1 = null;
	private byte[] r = null;        // random value
	private byte[] gr = null;       // digest of r
	private ByteBuffer messageBuffer;

	public EEncrypt() {
	}

	/**
	 * Create an elliptic key pair.
	 *
	 * @param    E      the elliptic curve containing the public key.
	 * @param    prng   a cryptographically strong pseudo-random number generator.
	 *
	 * @return the generated key pair (BigInteger x, EPoint Y),
	 * where x is the private key (0 < x < E.n) and Y is the normalized public key.
	 */
	public static Object[] makeKeyPair(ECurve E, SecureRandom prng) {
		BigInteger x;
		// generate suitable private key x:
		int t = E.n.bitLength();
		do {
			x = new BigInteger(t, prng);
			if (x.compareTo(E.n) >= 0) {
				x = x.subtract(E.n);
			}
			/* invariant: 0 <= x < n */
		} while (x.signum() == 0);
		/* invariant: 0 < x < n */
		// compute the corresponding public key:
		EPoint Y = E.kG(x).normalize();
		return new Object[]{x, Y};
	}

	/**
	 * Initialize the algorithm for encryption, with a given elliptic curve,
	 * given a suitable hash function and a prng.
	 *
	 * @param   E       elliptic curve.
	 * @param   md      hash function.
	 * @param   prng    pseudo-random number generator.
	 *
	 */
	void init(ECurve E, MessageDigest md, SecureRandom prng) {
		this.E = E;
		this.md = md;
		this.messageBuffer = new ByteBuffer();

		int len = (E.n.bitLength() + 7) / 8;
		this.r = new byte[len];
		prng.nextBytes(r);

		//Computes gr = hash(r) and resets
		this.gr = md.digest(r);
		this.md.reset();

		//Encrypting data
		this.state = STATE_ENCRYPTING;
	}

	/**
	 * Initialize the encryption algorithm for decryption,
	 * with a given elliptic curve, given a suitable hash function
	 * and the ciphertext parts C1 e c2.
	 *
	 * @param   E       elliptic curve.
	 * @param   md      hash function.
	 *
	 */
	void init(ECurve E, MessageDigest md, BigInteger privKey, EPoint C1, byte[] c2) {
		this.E = E;
		this.md = md;
		this.messageBuffer = new ByteBuffer();
		this.C1 = C1;

		EPoint D = C1.multiply(privKey).normalize();

		int len = (E.n.bitLength() + 7) / 8;
		this.r = new byte[len];

		//Compute r = c2 XOR x_d
		byte[] normalizedXd = new byte[r.length];
		appendFixedBytes(D.x.toByteArray(), normalizedXd, 0, r.length);
		for (int i = 0; i < r.length; i++) {
			r[i] = (byte) (c2[i] ^ normalizedXd[i]);
		}

		//Computes gr = hash(r) and resets
		this.gr = md.digest(r);
		this.md.reset();

		//Encrypting data
		this.state = STATE_DECRYPTING;
	}

	void update(byte[] msg) {
		this.messageBuffer.add(msg);
		if (this.state == STATE_ENCRYPTING) {
			md.update(msg);
		}
	}

	private byte[] symmetricEncryption(byte[] data, byte[] key){

		byte[] ciphertext = new byte[data.length];

		for (int i = 0; i < data.length; i++) {
			ciphertext[i] = (byte) (data[i] ^ key[i]);
		}
/*
		byte[] truncatedKey = new byte[CIPHER_KEY_LENGTH];
		System.arraycopy(key, 0, truncatedKey, 0, truncatedKey.length);

		SecretKeySpec keySpec = new SecretKeySpec(truncatedKey, "AES");
		Cipher cipher = Cipher.getInstance("AES");

		// encryption pass
		cipher.init(Cipher.ENCRYPT_MODE, keySpec);

		byte[] ciphertext = new byte[cipher.getOutputSize(data.length)];
		int ctLength = cipher.update(data, 0, data.length, ciphertext, 0);
		ctLength += cipher.doFinal(ciphertext, ctLength);
 */
		return ciphertext;
	}

	private byte[] symmetricDecryption(byte[] data, byte[] key){

		return this.symmetricEncryption(data, key);
		/*
		byte[] truncatedKey = new byte[CIPHER_KEY_LENGTH];
		System.arraycopy(key, 0, truncatedKey, 0, truncatedKey.length);

		SecretKeySpec keySpec = new SecretKeySpec(truncatedKey, "AES");
		Cipher cipher = Cipher.getInstance("AES");

		// decryption pass
		cipher.init(Cipher.DECRYPT_MODE, keySpec);
		byte[] plainText = new byte[cipher.getOutputSize(data.length)];
		int ptLength = cipher.update(data, 0, data.length, plainText, 0);
		ptLength += cipher.doFinal(plainText, ptLength);
		return plainText;
		//*/
	}

	public Object[] encrypt(EPoint pubKey){
		//Finishes hash computation
		BigInteger t = new BigInteger(md.digest(r));

		EPoint Q = pubKey.multiply(t).normalize();
		EPoint R = E.G.multiply(t).normalize();

		//Compute c2 = r XOR x_q
		byte[] normalizedXq = new byte[r.length];
		byte[] c2 = new byte[r.length];
		appendFixedBytes(Q.x.toByteArray(), normalizedXq, 0, r.length);
		for (int i = 0; i < c2.length; i++) {
			c2[i] = (byte) (r[i] ^ normalizedXq[i]);
		}

		//Compute c3 = SymEnc(message, key)
		byte[] c3 = this.messageBuffer.toByteArray();
		c3 = symmetricEncryption(c3, this.gr);

		//Updates state: reinitialization needed
		this.state = STATE_NOT_INITIALIZED;

		Object[] obj = {R, c2, c3};
		return obj;
	}

	/**
	 *
	 * @return plaintext or <code>null</code> if plaintext invalid
	 */
	public byte[] decrypt(){

		byte[] c3 = this.messageBuffer.toByteArray();

		//Compute c3 = SymEnc(message, key)
		c3 = symmetricDecryption(c3, this.gr);

		//Updates the hash
		this.md.update(c3);
		byte[] h = this.md.digest(this.r);

		EPoint newC1 = E.G.multiply(new BigInteger(h)).normalize();

		//Updates state: reinitialization needed
		this.state = STATE_NOT_INITIALIZED;

		if (newC1.equals(C1)) {
			return c3;
		}

		return null;
	}

	private void appendFixedBytes(byte[] data, byte[] dst, int off, int len) {
		if (data.length <= len) {
			Arrays.fill(dst, 0, len - data.length, (byte) 0);
			System.arraycopy(data, 0, dst,
					off + len - data.length, data.length);
		} else {
			/*
			 * this can only happen if the first octet is negative
			 * (meaning a negative BigInteger in two's complement),
			 * in which case a zero octet is prepended to the byte array:
			 */
			assert (data.length == len + 1 && data[0] == 0);
			System.arraycopy(data, 1, dst, off, len);
		}
	}

	public static void main(String[] args) {
		int numBits = 160;
		String message = "test123";
		ECurve curve = ECurve.getStandardCurve(numBits); // load curve equation for E
		SecureRandom rand = new SecureRandom();

		System.out.print(message + ": ");
		System.out.println(Util.getHexString(message.getBytes()));

		Object[] pair = ESigns.makeKeyPair(curve, rand);
		if (pair == null) {
			System.out.println("unable to make key pair!");
			return;
		}
		BigInteger privKey = (BigInteger) pair[0];
		EPoint pubKey = (EPoint) pair[1];

		EEncrypt ee = new EEncrypt();
		Object[] ciphertext = null;

		try {
			ee.init(curve, MessageDigest.getInstance("SHA"), rand);
			ee.update(message.getBytes()); // signs the whole msg
			//System.out.print("u");
			ciphertext = ee.encrypt(pubKey);
			//System.out.print("s")

			EPoint C1 = (EPoint) ciphertext[0];
			byte[] c2 = (byte[]) ciphertext[1];
			byte[] c3 = (byte[]) ciphertext[2];

			ee.init(curve, MessageDigest.getInstance("SHA"),
					privKey, C1, c2);
			//System.out.print("i");
			//sv.update(msg, 0, i); // verify the first i characters of msg
			ee.update(c3); // verify the whole msg
			//System.out.print("u");

			System.out.print(new String(c3) + ": ");
			System.out.println(Util.getHexString(c3));

			byte[] plaintext = ee.decrypt();

			System.out.print(new String(plaintext) + ": ");
			System.out.println(Util.getHexString(message.getBytes()));

		} catch (Exception ex) {
			ex.printStackTrace();
		}
	}
}
